from scalpel.cfg import CFGBuilder
from scalpel.SSA.const import SSA
import json
import os
from gremlin_python import statics
from gremlin_python.process.anonymous_traversal import traversal
from gremlin_python.process.graph_traversal import __
from gremlin_python.process.strategies import *
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from gremlin_python.structure.graph import Graph
from gremlin_python.process.traversal import T
from gremlin_python.process.traversal import Order
from gremlin_python.process.traversal import Cardinality
from gremlin_python.process.traversal import Column
from gremlin_python.process.traversal import Direction
from gremlin_python.process.traversal import Operator
from gremlin_python.process.traversal import P
from gremlin_python.process.traversal import TextP
from gremlin_python.process.traversal import Pop
from gremlin_python.process.traversal import Scope
from gremlin_python.process.traversal import Barrier
from gremlin_python.process.traversal import Bindings
from gremlin_python.process.traversal import WithOptions
import ast
import astor

visited_block_first = {}
visited_block_last = {}


def to_src_without_children(node):
    if isinstance(node, ast.FunctionDef):
        node = ast.copy_location(
            ast.FunctionDef(name=node.name, args=node.args, body=[], decorator_list=node.decorator_list), node)
    elif isinstance(node, ast.ClassDef):
        node = ast.copy_location(ast.ClassDef(name=node.name, bases=node.bases, keywords=node.keywords, body=[],
                                              decorator_list=node.decorator_list), node)
    elif isinstance(node, ast.If):
        node = ast.copy_location(ast.If(test=node.test, body=[], orelse=[]), node)
    elif isinstance(node, ast.While):
        node = ast.copy_location(ast.While(test=node.test, body=[], orelse=[]), node)
    elif isinstance(node, ast.For):
        node = ast.copy_location(ast.For(target=node.target, iter=node.iter, body=[], orelse=[]), node)
    elif isinstance(node, ast.Try):
        node = ast.copy_location(ast.Try(body=[], handlers=[], orelse=[], finalbody=[]), node)
    elif isinstance(node, ast.With):
        node = ast.copy_location(ast.With(items=node.items, body=[]), node)

    # Append the line number to the source code

    return node, astor.to_source(node).strip().replace("\n", ""), node.lineno


def visit_blocks(graph, block, visited, calls):
    if block.id in visited:
        return

    visited.append(block.id)
    visited_block_first[block.id] = block.statements[0]
    visited_block_last[block.id] = block.statements[-1]

    # Get source code from block statements
    src_code = ''
    for stmt in block.statements:
        stmt_single, stmt_code, stmt_lineno = to_src_without_children(stmt)
        stmt_str = json.dumps(ast.dump(stmt_single))
        graph.addV("code").barrier().property("ast", stmt_str).barrier().property("lineo",
                                                                                  stmt_lineno).barrier().property(
            "code", stmt_code).barrier().next()
        result = graph.V().count().next()
        print(result)
        # 此处存入相应的语句
    for exit in block.exits:
        if calls or exit.name not in block.cfg.functioncfgs:
            visit_blocks(graph, exit.target, visited, calls)
    # Add source code to block label


#    for exit in block.exits:
#        if calls or exit.name not in block.cfg.functioncfgs:
#            visit_blocks(graph, exit.target, visited, calls)
#            graph.edge(str(block.id), str(exit.target.id),
#                       label=str(exit.exitcase) if exit.exitcase else None)


def add_edge(g, block, visited, calls):
    if block.id in visited:
        return
    visited.append(block.id)
    if len(block.statements) == 1:
        if block.exits:
            stmt = block.statements[0]
            stmt_single, stmt_code, stmt_lineno = to_src_without_children(stmt)
            stmt_str = json.dumps(ast.dump(stmt_single))
            from_v = g.V().has('ast', stmt_str).has('lineo', stmt_lineno).has('code', stmt_code).limit(1).next()
        else:
            return 0

        for exit in block.exits:
            if exit.target.id in visited:
                stmt_to = visited_block_last[exit.target.id]
                stmt_single_to, stmt_code_to, stmt_lineno_to = to_src_without_children(stmt_to)
                stmt_str_to = json.dumps(ast.dump(stmt_single_to))
                to_v = g.V().has('ast', stmt_str_to).has('lineo', stmt_lineno_to).has('code', stmt_code_to).limit(
                    1).next()
            else:
                stmt_to = visited_block_first[exit.target.id]
                stmt_single_to, stmt_code_to, stmt_lineno_to = to_src_without_children(stmt_to)
                stmt_str_to = json.dumps(ast.dump(stmt_single_to))
                to_v = g.V().has('ast', stmt_str_to).has('lineo', stmt_lineno_to).has('code', stmt_code_to).limit(
                    1).next()

            if (from_v != None and to_v != None):
                g.V(from_v).addE("cfg").to(to_v).next()
                result = g.E().count().next()
                print(result)
            else:
                print("节点查询有问题")

    else:
        for i in range(0, len(block.statements)):
            if i >= len(block.statements) - 1:
                if block.exits:
                    stmt = block.statements[i]
                    stmt_single, stmt_code, stmt_lineno = to_src_without_children(stmt)
                    stmt_str = json.dumps(ast.dump(stmt_single))
                    from_v = g.V().has('ast', stmt_str).has('lineo', stmt_lineno).has('code', stmt_code).limit(1).next()

                for exit in block.exits:
                    if exit.target.id in visited:
                        stmt_to = visited_block_last[exit.target.id]
                        stmt_single_to, stmt_code_to, stmt_lineno_to = to_src_without_children(stmt_to)
                        stmt_str_to = json.dumps(ast.dump(stmt_single_to))
                        to_v = g.V().hasLabel('code').has('ast', stmt_str_to).has('lineo', stmt_lineno_to).has('code',
                                                                                                               stmt_code_to).limit(
                            1).next()
                    else:
                        stmt_to = visited_block_first[exit.target.id]
                        stmt_single_to, stmt_code_to, stmt_lineno_to = to_src_without_children(stmt_to)
                        stmt_str_to = json.dumps(ast.dump(stmt_single_to))
                        to_v = g.V().hasLabel('code').has('ast', stmt_str_to).has('lineo', stmt_lineno_to).has('code',
                                                                                                               stmt_code_to).limit(
                            1).next()

                    if (from_v != None and to_v != None):
                        g.V(from_v).addE("cfg").to(to_v).next()
                        result = g.E().count().next()
                        print(result)
                    else:
                        print("节点查询有问题")
            else:
                stmt = block.statements[i]
                stmt_single, stmt_code, stmt_lineno = to_src_without_children(stmt)
                stmt_str = json.dumps(ast.dump(stmt_single))
                from_v = g.V().hasLabel('code').has('ast', stmt_str).has('lineo', stmt_lineno).has('code',
                                                                                                   stmt_code).limit(
                    1).next()
                stmt_to = block.statements[i + 1]
                stmt_single_to, stmt_code_to, stmt_lineno_to = to_src_without_children(stmt_to)
                stmt_str_to = json.dumps(ast.dump(stmt_single_to))
                to_v = g.V().hasLabel('code').has('ast', stmt_str_to).has('lineo', stmt_lineno_to).has('code',
                                                                                                       stmt_code_to).limit(
                    1).next()
                if (from_v != None and to_v != None):
                    g.V(from_v).addE("cfg").to(to_v).next()
                    result = g.E().count().next()
                    print(result)
                else:
                    print("节点查询有问题")
    for exit in block.exits:
        if calls or exit.name not in block.cfg.functioncfgs:
            add_edge(g, exit.target, visited, calls)


def get_all_block_of_cfgs(flattened_cfg):
    all_block = []
    for fqn, cfg in flattened_cfg.items():
        all_block += cfg.get_all_blocks()

    return all_block


def build_and_input_cfgs(flattened_cfg, g):
    for fqn, cfg in flattened_cfg.items():
        visit_blocks(g, cfg.entryblock, visited=[], calls=True)
        add_edge(g, cfg.entryblock, visited=[], calls=True)


def get_all_block_of_cfgs_dict(all_block):
    return {block.id: block for block in all_block}


def scalpel_cfg_build(test_scripts):
    for script_path in test_scripts:
        test_script = os.path.basename(script_path)
        cfg_builder = CFGBuilder()
        flattened_cfg = cfg_builder.build_from_file(test_script,
                                                    script_path,
                                                    flattened=True)
        all_block = get_all_block_of_cfgs(flattened_cfg)
        all_block_dict = get_all_block_of_cfgs_dict(all_block)
        g = Graph().traversal().withRemote(DriverRemoteConnection('ws://localhost:8182/gremlin', 'g'))
        g.V().drop().iterate()
        build_and_input_cfgs(flattened_cfg, g)


if __name__ == "__main__":
    test_scripts = [r"/Research-Classroom/GreNeoGraphQL:A vulnerability detection DSL based on code graph database queries./example/Code_to_be_tested.py"]#修改为相应的位置
    scalpel_cfg_build(test_scripts)

